from tqdm import tqdm
from copy import deepcopy

from VRP_SA.parameter import *
from VRP_SA.GeneClass import *

DEBUG = 0

# return a bunch of random genes
def getRandomGenes(size):
    genes = []
    i = 0
    while i < size:
        gene = Gene("Gene " + str(i))
        if gene.isFeasible():
            genes.append(gene)
            i += 1
    return genes


# 更新选择概率
def updateChooseProb(genes):
    sumFit = 0
    for gene in genes:
        sumFit += gene.fit
    for gene in genes:
        gene.updateChooseProb(sumFit)


# 计算累计概率
def getSumProb(genes):
    sumProb = 0
    for gene in genes:
        sumProb += gene.chooseProb
    return sumProb


# 选择复制，选择前 1/3
def choose(genes):
    num = int(geneNum / 6) * 2  # 选择偶数个，方便下一步交叉
    # sort genes with respect to chooseProb
    key = lambda gene: gene.chooseProb
    genes.sort(reverse=True, key=key)
    # return shuffled top 1/3
    return genes[0:num]


# Position-based Crossover
def crossPBX(data1, data2, randIndex):
    new_data = [None] * len(data1)

    nums = [data1[i] for i in randIndex]
    _data2 = copy.deepcopy(data2)
    # _data2 = data2
    for i in range(len(randIndex)):
        new_data[randIndex[i]] = nums[i]
        _data2.remove(nums[i])
    for i in range(len(_data2)):
        for j, new_num in enumerate(new_data):
            if new_num is None:
                new_data[j] = _data2[i]
                break
    return new_data

# Partial-Mapped Crossover 随机选一段交换
def crossPMX(data1, data2, randIndex):
    i, j = randIndex
    while i <= j:
        t = data1[i]
        data1[i] = data2[i]
        data2[i] = t
        i += 1

# 交叉一对
def crossPair(parent1, parent2, crossedGenes):
    parent1_customer = parent1.readData()[0]
    parent2_customer = parent2.readData()[0]
    # 对首尾的0不做交叉
    parent1_customer.pop(0)
    parent1_customer.pop()
    parent2_customer.pop(0)
    parent2_customer.pop()

    length1 = len(parent1_customer)
    randIndex1 = random.sample(range(length1), (length1 // 6) * 2)
    offspring1_customer = [0] + crossPBX(parent1_customer, parent2_customer, randIndex1) + [0]
    offspring2_customer = [0] + crossPBX(parent2_customer, parent1_customer, randIndex1) + [0]
    length2 = len(parent1.readData()[1])
    randIndex2 = sorted(random.sample(range(length2), 2)) # 起止位置
    offspring1_bid = parent1.readData()[1]
    offspring2_bid = parent2.readData()[1]
    crossPMX(offspring1_bid, offspring2_bid, randIndex2)

    offspring1 = Gene(data=[offspring1_customer, offspring1_bid])
    offspring2 = Gene(data=[offspring2_customer, offspring2_bid])
    if offspring1.isFeasible() and offspring2.isFeasible():
        crossedGenes.append(offspring1)
        crossedGenes.append(offspring2)
        return True
    else:
        return False


# 交叉
def cross(genes):
    crossedGenes = []
    i = 0
    while i < len(genes):
        if crossPair(genes[i], genes[i + 1], crossedGenes):
            i += 2
    return crossedGenes


# 合并
def mergeGenes(genes, crossedGenes):
    # sort genes with respect to chooseProb
    key = lambda gene: gene.chooseProb
    genes.sort(reverse=True, key=key)
    pos = geneNum - 1
    for gene in crossedGenes:
        genes[pos] = gene
        pos -= 1
    return genes


# 变异
def vary(genes):
    for index, gene in enumerate(genes):
        # 精英主义，保留前三十
        if index < 30:
            continue
        if random.random() < VARY:
            genes[index] = varyOne(gene)
    return genes


def varyOne(gene):
    varyNum = 10
    variedGenes = []
    if random.random() < 0.5:  # 变异customer序列
        for _ in range(varyNum):
            i1, i2 = random.sample(range(1, len(gene.readData()[0]) - 1), k=2)
            new = gene.readData()[0]
            new[i1], new[i2] = new[i2], new[i1]
            variedGenes.append(Gene(data=[new.copy(), gene.readData()[1]]))
    else:  # 变异Bid序列
        for _ in range(varyNum):
            i = random.randint(0, len(gene.readData()[1])-1)
            new = gene.readData()[1]
            new[i] = 1 - new[i]
            variedGenes.append(Gene(data=[gene.readData()[0], new.copy()]))
    key = lambda gene: gene.fit
    variedGenes.sort(reverse=True, key=key)
    return variedGenes[0]


if __name__ == "__main__" and DEBUG == 0:
    genes = getRandomGenes(geneNum)  # 初始种群
    # 迭代
    for i in tqdm(range(generationNum)):
        updateChooseProb(genes)
        sumProb = getSumProb(genes)
        chosenGenes = choose(deepcopy(genes))  # 选择
        crossedGenes = cross(chosenGenes)  # 交叉
        genes = mergeGenes(genes, crossedGenes)  # 复制交叉至子代种群
        genes = vary(genes)  # 变异
    # sort genes with respect to chooseProb
    key = lambda gene: gene.fit
    genes.sort(reverse=True, key=key)  # 以fit对种群排序
    print('\r\n')
    print('objective:', 1/genes[0].fit)
    genes[0].plot()  # 画出来

if DEBUG == 1:
    a = [3, 15, 21, 14, 23, 16, 19, 2, 20, 0, 13, 11, 7, 4, 0, 12, 8, 5, 24, 22, 0, 10, 6, 18, 9, 1, 17]
    b = [23, 12, 8, 14, 4, 21, 15, 7, 24, 22, 0, 13, 6, 5, 17, 19, 0, 16, 3, 9, 10, 20, 11, 18, 0, 1, 2]

